#!/usr/bin/env python
# coding: utf-8

import jwt
import time
import requests
class LLM(object):
    def __init__(self, url, api_key):
        self.url = url
        self.api_key = api_key
        self.token = None

    # 调用大模型
    def get_llm(self, question):
        if self.token is None:
            self.token = self.__generate_token()

        headers = {
            'Content-Type': 'application/json',
            'Authorization': self.token
        }
        data = {
            "model": "glm-4",
            "messages": [{
                "role": "user",
                "content": question
            }]}
        response = requests.post(self.url, headers=headers, json=data)
        return response.json()
    # 获取glm的token
    def __generate_token(self, exp_seconds: int=3600):
        try:
            id, secret = self.api_key.split(".")
        except Exception as e:
            raise Exception("invalid apikey", e)

        payload = {
            "api_key": id,
            "exp": int(round(time.time() * 1000)) + exp_seconds * 1000,
            "timestamp": int(round(time.time() * 1000)),
        }
        return jwt.encode(
            payload,
            secret,
            algorithm="HS256",
            headers={"alg": "HS256", "sign_type": "SIGN"},
        )
